#include <ntddk.h>
#include <wdm.h>

typedef unsigned int UINT;
typedef char * PCHAR;

typedef struct tagDEVICE_EXTENSION {
    PVOID DataBuffer;                  // Internal Buffer for Read/Write I/O
    ULONG DataBufferSize;
} DEVICE_EXTENSION, *PDEVICE_EXTENSION;

NTSTATUS Close(PDEVICE_OBJECT DeviceObject, PIRP Irp);
NTSTATUS WriteBufferedIO(PDEVICE_OBJECT DeviceObject, PIRP Irp);
NTSTATUS ReadBufferedIO(PDEVICE_OBJECT DeviceObject, PIRP Irp);
NTSTATUS UnSupportedFunction(PDEVICE_OBJECT DeviceObject, PIRP Irp);
VOID Unload(PDRIVER_OBJECT  DriverObject); 

NTSTATUS DriverEntry(PDRIVER_OBJECT  pDriverObject, PUNICODE_STRING  pRegistryPath)
{
    NTSTATUS NtStatus = STATUS_SUCCESS;
    UINT uiIndex = 0;
    PDEVICE_OBJECT pDeviceObject = NULL;
    UNICODE_STRING usDriverName, usDosDeviceName;
    PDEVICE_EXTENSION p_DVCEXT;
    ULONG DeviceExtensionSize;

    DeviceExtensionSize = sizeof(DEVICE_EXTENSION);

    DbgPrint("DriverEntry Called \r\n");

    RtlInitUnicodeString(&usDriverName, L"\\Device\\MyDevice");
    RtlInitUnicodeString(&usDosDeviceName, L"\\DosDevices\\MyDevice"); 

    NtStatus = IoCreateDevice(pDriverObject, DeviceExtensionSize, &usDriverName, FILE_DEVICE_UNKNOWN, FILE_DEVICE_SECURE_OPEN, FALSE, &pDeviceObject);

    if(NtStatus == STATUS_SUCCESS)
    {

        for(uiIndex = 0; uiIndex < IRP_MJ_MAXIMUM_FUNCTION; uiIndex++)
             pDriverObject->MajorFunction[uiIndex] = UnSupportedFunction;
    
        pDriverObject->MajorFunction[IRP_MJ_CLOSE]             = Close;
        pDriverObject->MajorFunction[IRP_MJ_READ]              = ReadBufferedIO;
        pDriverObject->MajorFunction[IRP_MJ_WRITE]             = WriteBufferedIO;
    
        pDriverObject->DriverUnload = Unload;

        pDeviceObject->Flags |= DO_BUFFERED_IO;
    
        pDeviceObject->Flags &= (~DO_DEVICE_INITIALIZING);

        p_DVCEXT = pDeviceObject->DeviceExtension;
        p_DVCEXT->DataBuffer = NULL;
        p_DVCEXT->DataBufferSize = 0;
    
        IoCreateSymbolicLink(&usDosDeviceName, &usDriverName);
    }


    return NtStatus;
}

VOID Unload(PDRIVER_OBJECT  DriverObject)
{    
    
    UNICODE_STRING usDosDeviceName;
    
    DbgPrint("Unload Called \r\n");
    
    RtlInitUnicodeString(&usDosDeviceName, L"\\DosDevices\\MyDevice");
    IoDeleteSymbolicLink(&usDosDeviceName);

    IoDeleteDevice(DriverObject->DeviceObject);
}

NTSTATUS Close(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
    NTSTATUS NtStatus = STATUS_SUCCESS;
    DbgPrint("Close Called \r\n");

    return NtStatus;
}

BOOLEAN IsStringTerminated(PCHAR pString, UINT uiLength, UINT *pdwStringLength)
{
    BOOLEAN bStringIsTerminated = FALSE;
    UINT uiIndex = 0;
    
    DbgPrint("IsStringTerminated(0x%0x, %d)\r\n", pString, uiLength);

    *pdwStringLength = 0;

    while(uiIndex < uiLength && bStringIsTerminated == FALSE)
    {
        if(pString[uiIndex] == '\0')
        {
            *pdwStringLength = uiIndex + 1; /* Include the total count we read, includes the NULL */
            bStringIsTerminated = TRUE;
            DbgPrint("  String Is Terminated!\r\n");
        }
        else
        {
           uiIndex++;
        }
    }

    return bStringIsTerminated;
}

NTSTATUS WriteBufferedIO(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
    NTSTATUS NtStatus = STATUS_UNSUCCESSFUL;
    PIO_STACK_LOCATION pIoStackIrp = NULL;
    UINT dwDataWritten = 0;
    PCHAR pWriteDataBuffer;
    ULONG messageSize = 0;

    PDEVICE_EXTENSION p_DVCEXT = (PDEVICE_EXTENSION)DeviceObject->DeviceExtension;
    if (p_DVCEXT->DataBuffer != NULL) {
        ExFreePool(p_DVCEXT->DataBuffer);
        p_DVCEXT->DataBuffer = NULL;
        p_DVCEXT->DataBufferSize = 0;
    }

    DbgPrint("WriteBufferedIO Called \r\n");
    
    pIoStackIrp = IoGetCurrentIrpStackLocation(Irp);
    
    if(pIoStackIrp)
    {
        messageSize = pIoStackIrp->Parameters.Write.Length;
        pWriteDataBuffer = (PCHAR)Irp->AssociatedIrp.SystemBuffer;
    
        if(pWriteDataBuffer)
        {     
            p_DVCEXT->DataBuffer = ExAllocatePool( NonPagedPool, messageSize );                        
            if(IsStringTerminated(pWriteDataBuffer, pIoStackIrp->Parameters.Write.Length, &dwDataWritten) && p_DVCEXT->DataBuffer != NULL)
            {
                DbgPrint(pWriteDataBuffer);
                p_DVCEXT->DataBufferSize = messageSize;
                RtlCopyMemory( p_DVCEXT->DataBuffer, pWriteDataBuffer, messageSize );
                NtStatus = STATUS_SUCCESS;
            }
        }
    }

    Irp->IoStatus.Status = NtStatus;
    Irp->IoStatus.Information = messageSize;

    IoCompleteRequest(Irp, IO_NO_INCREMENT);


    return NtStatus;
}

NTSTATUS ReadBufferedIO(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
    NTSTATUS NtStatus = STATUS_BUFFER_TOO_SMALL;
    PIO_STACK_LOCATION pIoStackIrp = NULL;
    PCHAR pReadDataBuffer;
    ULONG messageSize = 0;

    PDEVICE_EXTENSION p_DVCEXT = (PDEVICE_EXTENSION)DeviceObject->DeviceExtension;

    DbgPrint("ReadBufferedIO Called \r\n");
    
    pIoStackIrp = IoGetCurrentIrpStackLocation(Irp);
    
    if(pIoStackIrp)
    {
        messageSize = pIoStackIrp->Parameters.Read.Length;
        pReadDataBuffer = (PCHAR)Irp->AssociatedIrp.SystemBuffer;

        if ( messageSize > p_DVCEXT->DataBufferSize) {
            messageSize = p_DVCEXT->DataBufferSize;
        }
    
        if(pReadDataBuffer)
        {                             
            RtlCopyMemory(pReadDataBuffer, p_DVCEXT->DataBuffer, messageSize);
            NtStatus = STATUS_SUCCESS;
        }
    }

    ExFreePool(p_DVCEXT->DataBuffer);
    p_DVCEXT->DataBuffer = NULL;
    p_DVCEXT->DataBufferSize = 0;

    Irp->IoStatus.Status = NtStatus;
    Irp->IoStatus.Information = messageSize;

    IoCompleteRequest(Irp, IO_NO_INCREMENT);

    return NtStatus;
}

NTSTATUS UnSupportedFunction(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
    NTSTATUS NtStatus = STATUS_NOT_SUPPORTED;
    DbgPrint("UnSupportedFunction Called \r\n");

    return NtStatus;
}